# pip install markdown-it

from markdown_it import MarkdownIt
import os, re, json

def parse_md(md_file):
    with open(md_file, 'r', encoding='utf-8') as f:
        content = f.read()
    md = MarkdownIt().enable("table")
    tokens = md.parse(content)
    
    return tokens

def filter_content(content):
    filtered_text = re.sub(r"!\[\]\(.*?\)", "", content)
    
    if content.find("<") >= 0:
        filtered_text = re.sub(r'<span[^>]*></span>', '', filtered_text)
        
    return filtered_text

def json_to_md_table(json_data):
    # 解析JSON数据
    header = json_data["headers"]
    table_content = json_data["rows"]

    # 构建Markdown表格
    md_table = ""

    # 添加表头
    md_table += "| " + " | ".join(header) + " |\n"
    md_table += "| " + " | ".join(["---" for _ in header]) + " |\n"

    # 添加表格内容
    for row in table_content:
        md_table += "| " + " | ".join(map(str, row)) + " |\n"

    return md_table

# 示例JSON数据
json_data = {
    "header": ["Name", "Age", "City"],
    "table_content": [
        ["Alice", 25, "New York"],
        ["Bob", 30, "Los Angeles"],
        ["Charlie", 35, "Chicago"]
    ]
}

def extract_one_table(tokens):
    current_table = []
    current_row = []
    current_cell = []
    in_table = False
    in_thead = False
    in_tbody = False

    for idx, token in enumerate(tokens):
        if token.type == "table_open":
            in_table = True
            current_table = { "headers": [], "rows": [] }
        elif token.type == "table_close":
            in_table = False
            return idx, current_table
        
        elif token.type == "thead_open":
            in_thead = True
        elif token.type == "thead_close":
            in_thead = False
        elif token.type == "tbody_open":
            in_tbody = True
        elif token.type == "tbody_close":
            in_tbody = False
        elif token.type == "tr_open":
            current_row = []
        elif token.type == "tr_close":
            if in_thead:
                current_table["headers"] = current_row
            elif in_tbody:
                current_table["rows"].append(current_row)
        elif token.type == "th_open" or token.type == "td_open":
            current_cell = []
        elif token.type == "th_close" or token.type == "td_close":
            cell_content = "".join(current_cell).strip()
            current_row.append(cell_content)
            current_cell = []
        elif token.type == "inline":
            current_cell.append(token.content)

    return -1, current_table


def extract_tables(tokens):
    tables = []
    current_table = []
    current_row = []
    current_cell = []
    in_table = False
    in_thead = False
    in_tbody = False

    for token in tokens:
        if token.type == "table_open":
            in_table = True
            current_table = { "headers": [], "rows": [] }
        elif token.type == "table_close":
            in_table = False
            tables.append(current_table)
        elif token.type == "thead_open":
            in_thead = True
        elif token.type == "thead_close":
            in_thead = False
        elif token.type == "tbody_open":
            in_tbody = True
        elif token.type == "tbody_close":
            in_tbody = False
        elif token.type == "tr_open":
            current_row = []
        elif token.type == "tr_close":
            if in_thead:
                current_table["headers"] = current_row
            elif in_tbody:
                current_table["rows"].append(current_row)
        elif token.type == "th_open" or token.type == "td_open":
            current_cell = []
        elif token.type == "th_close" or token.type == "td_close":
            cell_content = "".join(current_cell).strip()
            current_row.append(cell_content)
            current_cell = []
        elif token.type == "inline":
            current_cell.append(token.content)

    return tables

def extract_json_contents(tokens, company="", max_header_length=64):
    '''
    json format:
    company: 公司名称
    headers: 多级headers拼接的串
    source: default or vl
    content: 文本内容
    '''
    
    json_contents = []
    contexts = {}
    in_heading = False
    heading_level = ""
    in_paragraph = False
    paragraph_contents = []
    
    idx = 0
    while idx < len(tokens):
        token = tokens[idx]
        
        if token.type == "table_open":
            offset, table = extract_one_table(tokens[idx:])
            
            if offset >= 0:
                idx += offset
                
                headers = []
                for i in range(1, 10):
                    key = "h"+str(i)
                    if key in contexts:
                        headers.append(contexts[key])
                        
                item = {}
                item["company"] = company 
                item["headers"] = "#".join(headers)
                item['source'] = "default"
                item["content"] = json_to_md_table(table)
                json_contents.append(item)
                
            idx += 1
            continue    
        
        if token.type == 'heading_open':
            in_heading = True 
            heading_level = token.tag
            # 清空低层级的标题
            level = int(heading_level[1:])
            for i in range(level, 10):
                key = "h"+str(i)
                if key in contexts:
                    contexts.pop(key)
                    
        elif token.type == 'heading_close':
            in_heading = False
            if heading_level in contexts and len(contexts[heading_level]) > max_header_length:
                paragraph_contents.append(contexts[heading_level])
                contexts.pop(heading_level)
                
                headers = []
                for i in range(1, 10):
                    key = "h"+str(i)
                    if key in contexts:
                        headers.append(contexts[key])
            
                item = {}
                item["company"] = company 
                item["headers"] = "#".join(headers)
                item['source'] = "default"
                item["content"] = "\n".join(paragraph_contents)
                json_contents.append(item)

                paragraph_contents = []
                
            heading_level = ""
            
        elif token.type == 'paragraph_open':
            in_paragraph = True
        elif token.type == 'paragraph_close':
            in_paragraph = False
            if len(paragraph_contents) == 0:
                idx += 1
                continue
            
            headers = []
            for i in range(1, 10):
                key = "h"+str(i)
                if key in contexts:
                    headers.append(contexts[key])
            
            item = {}
            item["company"] = company 
            item["headers"] = "#".join(headers)
            item['source'] = "default"
            item["content"] = "\n".join(paragraph_contents)
            json_contents.append(item)
            
            paragraph_contents = []
            
        elif token.type == 'inline':
            content = filter_content(token.content.strip())
            #if content.find("2019 年 4 月 26 日") >= 0:
            #    pass
            if content == "":
                idx += 1
                continue
            if in_heading:
                contexts[heading_level] = content
            if in_paragraph:
                paragraph_contents.append(content)
        idx += 1
                
    return json_contents

def combine_and_split(json_contents):
    new_json_contents = []
    tracking_headers = ""
    tracking_source = ""
    tracking_company = ""
    
    paragraph_contents = []
    for item in json_contents:
        #if item['content'].find("由锂科有限于 2019 年 4 月整体变更设立股份有限公司") >= 0:
        #    pass
        
        if item["headers"] == tracking_headers and item["source"] == tracking_source and item["company"] == tracking_company:
            paragraph_contents.append(item["content"])
        else:
            if len(paragraph_contents) > 0:
                new_item = {}
                new_item['company'] = tracking_company
                new_item['headers'] = tracking_headers
                new_item['source'] = tracking_source
                new_item["content"] = "\n".join(paragraph_contents)
                new_json_contents.append(new_item)
                
            paragraph_contents = [item["content"]]
            tracking_headers = item["headers"]
            tracking_source = item["source"]
            tracking_company = item["company"]
         
    if len(paragraph_contents) > 0:
        new_item = {}
        new_item['company'] = tracking_company
        new_item['headers'] = tracking_headers
        new_item['source'] = tracking_source
        new_item["content"] = "\n".join(paragraph_contents)
        new_json_contents.append(new_item)
        
        paragraph_contents = []
        tracking_headers = item["headers"]
        tracking_source = item["source"]
        tracking_company = item["company"]
        
    return new_json_contents

def extract_text_contents(md_path, company=""):
    tokens = parse_md(md_path)
    json_contents = extract_json_contents(tokens)
    json_contents = combine_and_split(json_contents)
    
    if company:
        for item in json_contents:
            item['company'] = company
            if item['headers'].find(company) < 0:
                item['headers'] = company + "#" + item['headers']

    return json_contents

# main
if __name__ == '__main__':
    md_file = "/Users/danningjiang/work/teaching/rag_projects/data_process/pdfout/eec23035376ae0e339a7643402fdbdccd92ad703/eec23035376ae0e339a7643402fdbdccd92ad703.md"
    company = "湖南长远锂科股份有限公司"
    json_contents = extract_text_contents(md_file, company=company)
    
    #tokens = parse_md(md_file)
    #json_contents = extract_json_contents(tokens)
    #json_contents = combine_and_split(json_contents)
    
    output_file = company + ".jsonl"
    print(output_file)
    
    with open(output_file, 'w', encoding='utf-8') as f:
        for item in json_contents:
            f.write(json.dumps(item, ensure_ascii=False)+"\n")
            
            